package com.countertack.endtoendstories;

import java.io.File;
import java.io.FileOutputStream;
import java.net.InetAddress;
import java.sql.ResultSet;
import java.util.UUID;

import com.countertack.impala.ImpalaConnector;
import com.countertack.ssh.SSHManager;
import com.countertack.utils.Configuration;
import com.countertack.utils.HttpClientUtils;
import com.countertack.vmware.VMWareConnector;
import com.countertack.vmware.VMWareUtils;
import com.vmware.vim25.mo.VirtualMachine;

public class ImpactModel {
	static String vmUser = null;
	static String vmPassword = null;

	public static void main(String[] args) throws Exception {

	
		/*  if (args.length <= 0) {
		  
		  System.out.println("configuration details missed !"); return; }
		  String configFilePath = args[0];*/
		

		String configFilePath = "C:\\Users\\user\\Desktop\\endtoendautomation\\config.properties";
		Configuration.loadProperties(configFilePath);
		// String configFilePath =
		// "C:\\Users\\user\\Desktop\\configm.properties";
		vmUser = Configuration.getProperty("vmuser");
		vmPassword = Configuration.getProperty("vmpassword");
		String uuid = UUID.randomUUID().toString();

		ImpalaConnector iConnector;

		iConnector = ImpalaConnector
				.createImpalaConnectorInstance(Configuration
						.getProperty("sentinelehost"));

		InetAddress address = InetAddress.getByName(Configuration
				.getProperty("sentinelehost"));
		String sentenelIp = address.getHostAddress();
		System.out.println(address.getHostAddress());
		// connect to esxi and starting virtual machines

		VMWareConnector esxiConnector = new VMWareConnector(
				Configuration.getProperty("esxi"));

		// poweredOn

		runMalwareTest(esxiConnector, iConnector, sentenelIp, "darkseoul.exe");
		System.out.println("execution completed");

	}

	public static void runMalwareTest(VMWareConnector esxiConnector,
			ImpalaConnector iConnector, String sentenelIp, String malwaresPath)
			throws Exception {
		// String endptIp = "10.140.62.159";
		File malwaref = new File(malwaresPath);
		try {
			if (!malwaref.exists()) {

				System.out.println("malware not exist");
				return;
			}
		} catch (Exception e) {
			System.out.println("malware not exist");
			e.printStackTrace();
			return;
		}

		final VirtualMachine endPointMachine = esxiConnector
				.getVirtualMachineByname(Configuration.getProperty("vmname"));

		String smokeSnapShotName = Configuration
				.getProperty("rootsnapshotname");
		esxiConnector.revertToSnapShort(endPointMachine.getSnapshot()
				.getRootSnapshotList(), smokeSnapShotName);

		System.out.println("Reverting to root snap shot>>>>>"
				+ smokeSnapShotName);

		Thread.sleep(180000);

		// win7x64-sym-2
		System.out.println("power state >>>"
				+ endPointMachine.getRuntime().powerState);

		String endpointMachineStatus = esxiConnector
				.getVirtualMachineStatus(endPointMachine);

		if (endpointMachineStatus == null) {

			System.out.println("Power on  endpoint machine .....");

			esxiConnector.startVirtualMachine(endPointMachine);

			Thread.sleep(600000);

		} else {

			if ("notRunning".equalsIgnoreCase(endpointMachineStatus)) {

				esxiConnector.reStartVirtualMachine(endPointMachine);

				System.out.println("Rebooting endpoint machine .....");

				Thread.sleep(600000);

			}

		}
boolean isXpmachine=false;
		
		String osType = VMWareUtils.getVMOSArchitecture(
				Configuration.getProperty("esxi"),
				Configuration.getProperty("vmname"));
		
		if (osType != null) {

			osType = osType.toLowerCase();

			if (osType.contains("xp")) {

				isXpmachine = true;
			}
		}


		// initialising ssh to virtual machines
		String endptIp = Configuration.getProperty("vmip");
		if (endptIp == null || endptIp.trim() == "") {

			endptIp = endPointMachine.getGuest().ipAddress;
		}

		System.out.println("ip>>>" + endptIp);

		SSHManager endpointSsh = new SSHManager(vmUser, vmPassword, endptIp);
		endpointSsh.connect();

		// install agent

		endpointSsh.copyFile(Configuration.getProperty("driverpath"),
				"CounterTack Agent Installer Hiding.exe");

		if (isXpmachine) {
			endpointSsh
					.runCommand("cmd /c \"C:\\Documents and Settings\\"
							+ vmUser
							+ "\\CounterTack Agent Installer Hiding.exe\" /S /HOST:"
							+ sentenelIp + " /PORT:443");

		} else {

			endpointSsh
					.runCommand("cmd /c \"C:\\Users\\"
							+ vmUser
							+ "\\CounterTack Agent Installer Hiding.exe\" /S /HOST:"
							+ sentenelIp + " /PORT:443");

		}


		System.out.println("Installing agent on endpoint");
		Thread.sleep(120000);

		String eId = null;
		ResultSet endpointIdRes = iConnector
				.executeQueryAndReturnResults("select id from endpoints where ip='"
						+ endptIp
						+ "'"
						+ " and status='ON' and ip is not null order by time_started desc limit 1");

		while (endpointIdRes.next()) {
			try {
				eId = endpointIdRes.getString(1);
			} catch (Exception e) {

			}

			System.out.println("endpoint Id>>>>>>" + eId);

		}

		if (eId == null) {

			System.out.println("endpoint is not installed");

			return;

		}

		pauseEndpoint(eId);
		Thread.sleep(60000);

		System.out.println("malware mname>>>" + malwaref.getName());

		String malwareTestString = malwaref.getName();
		if (malwaref.getName().length() > 4) {

			malwareTestString = malwaref.getName().substring(0, 4);

		}
		endpointSsh.copyFile(malwaresPath, malwaref.getName());

		Thread.sleep(6000);
		System.out.println("malware copied success fully");

		resumeEndpoint(eId);

		Thread.sleep(60000);

		try {
			if(isXpmachine){
				endpointSsh.runCommand("cmd /c \"C:\\Documents and Settings\\" + vmUser + "\\"
						+ malwaref.getName()+"\"");
				
			}else{
				
				
				endpointSsh.runCommand("cmd /c \"C:\\Users\\" + vmUser + "\\"
						+ malwaref.getName()+"\"");	
			}
			

			System.out.println("malware installed...");

		} catch (Exception e) {

		}
		boolean isbsod = false;
		endpointMachineStatus = esxiConnector
				.getVirtualMachineStatus(endPointMachine);
		if ("Running".equalsIgnoreCase(endpointMachineStatus)) {

			String bsodRes = endpointSsh
					.returnCommandOutPut("cmd /c cd C:\\Windows && dir /b /a-d");

			if (bsodRes != null) {

				if (bsodRes.contains("MEMORY.dump")) {

					isbsod = true;

				}

			}
		}
		Thread.sleep(1500000);

		// retriving behavior id

		ResultSet brs = iConnector
				.executeQueryAndReturnResults("select hex(id) from behaviors where uri like '%"
						+ eId + "%'  order by event_start_time desc limit 1");
		String bhid = null;
		boolean isMalwareTriggered = false;
		while (brs.next()) {
			bhid = brs.getString(1);

		}
		System.out.println("behavior id>>>>" + bhid);
		StringBuffer resultBuffer = new StringBuffer();
		resultBuffer.append("Sentinel>>>>" + sentenelIp + "\n");
		resultBuffer.append("VM Name>>>>>"
				+ Configuration.getProperty("vmname") + "\n");
		resultBuffer.append("VM IPAddd>>>>" + endptIp);
		resultBuffer.append("EndpointID>>>" + eId + "\n");
		String kernalModVer = getDriverVersion(eId, iConnector);
		resultBuffer.append("Kernal Module Version>>>" + kernalModVer + "\n");
		resultBuffer.append("Behavior id :::::" + bhid + "\n");
		int no_of_files = 0;
		int no_of_registries = 0;
		int no_of_processes = 0;
		int notification_count = 0;
		int key_event_count = 0;
		int event_count = 0;
		int behavior_count = 0;
		int basic_event_count = 0;

		if (bhid != null) {

			ResultSet ressProcess = iConnector
					.executeQueryAndReturnResults("select target_process_name from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			resultBuffer.append("<------------Processes Names-----------> \n");

			while (ressProcess.next()) {
				String processName = ressProcess.getString(1);

				if (processName != null) {

					if (processName.contains(malwareTestString)) {

						isMalwareTriggered = true;

					}
					resultBuffer.append(processName + "\n");
					no_of_processes++;
					// System.out.println("Process name>>>" + processName);

				}
			}

			ResultSet flwrrs = iConnector
					.executeQueryAndReturnResults("select target_file_path from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			resultBuffer.append("<------------File Names-----------> \n");
			while (flwrrs.next()) {
				String fn = flwrrs.getString(1);

				if (fn != null) {
					resultBuffer.append(fn + "\n");
					// System.out.println("file name>>>" + fn);
					no_of_files++;

				}
			}

			ResultSet regwrrs = iConnector
					.executeQueryAndReturnResults("select target_registry_key from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			// regkey
			resultBuffer.append("<------------Registry Keys-----------> \n");
			while (regwrrs.next()) {
				String rn = regwrrs.getString(1);

				if (rn != null) {
					resultBuffer.append(rn + "\n");
					// System.out.println("Reg Key>>>" + rn);
					no_of_registries++;

				}
			}

		}

		ResultSet clsrs = iConnector
				.executeQueryAndReturnResults("select classification_name,endpoint_id from classified_endpoints where endpoint_id='"
						+ eId + "'");
		boolean clsoc = false;
		String clsName = null;
		String clsGroup = "";
		while (clsrs.next()) {
			clsName = clsrs.getString(1);

			if (clsName != null) {
				System.out.println("classification_name>>>" + clsName);
				clsoc = true;
				clsGroup = clsGroup + "_" + clsName;
				// break;

			}
		}
		clsName = clsGroup;
		if (clsName.startsWith("_")) {

			clsName = clsName.replaceFirst("_", "");
		}
		if (clsoc) {

			// System.out.println("Malware test result -------->Pass");
		} else {

			// System.out.println("Malware test result ---------->Fail");

		}

		try {
			ResultSet countRs = iConnector
					.executeQueryAndReturnResults("select notification_count,key_event_count,event_count,behavior_count,basic_event_count from endpoints where id='"
							+ eId + "'");

			while (countRs.next()) {
				try {
					notification_count = countRs.getInt(1);
					key_event_count = countRs.getInt(2);
					event_count = countRs.getInt(3);
					behavior_count = countRs.getInt(4);
					basic_event_count = countRs.getInt(5);
				} catch (Exception e) {

					e.printStackTrace();

				}

			}

		} catch (Exception e1) {

			e1.printStackTrace();

		}

		// checking the threat level

		ResultSet threatRes = iConnector
				.executeQueryAndReturnResults("select threat from endpoints where id='"
						+ eId + "'" + " order by time_started desc limit 1");

		String threat = null;
		while (threatRes.next()) {
			threat = threatRes.getString(1);
			;
			try {
				if(threat==null){
					resultBuffer.append("Threat Level" + threat + "\n");
					System.out.println("Threat Level" + threat);
					resultBuffer.append("Test case failed"   + "\n");
					System.out.println("Test case failed");
					
				}
			else if (threat.equalsIgnoreCase("medium")
						|| threat.equalsIgnoreCase("critical")
						|| threat.equalsIgnoreCase("high")) {
					System.out.println("Threat Level" + threat);
					System.out.println("Test case Pass");
				} else {
					System.out.println("Test case failed");

				}
			} catch (Exception e) {
				e.printStackTrace();
			}

			// System.out.println("Threat level>>>>>>" + threat);

		}

		// checking the Impact level

		ResultSet impactRes = iConnector
				.executeQueryAndReturnResults("select impact_level from endpoints where id='"
						+ eId + "'" + "  order by time_started desc limit 1");

		String impact = null;

		while (impactRes.next()) {
			impact = impactRes.getString(1);
			;
			try {
				if(impact==null){
					System.out.println("Impact level>>>>>>" + impact);
					resultBuffer.append("Impact level>>>>>>" + impact + "\n");
					System.out.println("Test case failed");
					resultBuffer.append("Test case failed" + "\n");
					
				}
			else if (impact.equalsIgnoreCase("medium")
						|| impact.equalsIgnoreCase("critical")
						|| impact.equalsIgnoreCase("high")) {
					System.out.println("Impact level>>>>>>" + impact);
					System.out.println("Test case passed");
					resultBuffer.append("Test case passed" + "\n");
				} else {
					System.out.println("Test case failed");
					resultBuffer.append("Test case failed" + "\n");

				}
			} catch (Exception e) {
				e.printStackTrace();
			}

		}

		String malName = malwaref.getName();

		malName = malName.replaceAll("\\.", "_");

		File malResultFile = new File("Impact_model_malaware_result_"
				+ Configuration.getProperty("vmname") + "_" + malName + ".txt");

		malResultFile.createNewFile();
		try {
			FileOutputStream fout = new FileOutputStream(malResultFile);
			resultBuffer.append("Malware behavior status:" + isMalwareTriggered
					+ "\n");
			resultBuffer
					.append("Malware classification name:" + clsName + "\n");
			resultBuffer.append("BSOD:" + isbsod + "\n");

			resultBuffer.append("Number of Files:" + no_of_files + "\n");

			resultBuffer.append("Number of Regestries:" + no_of_registries
					+ "\n");

			resultBuffer
					.append("Number of Processes:" + no_of_processes + "\n");
			resultBuffer
					.append("Number of Keyevents:" + key_event_count + "\n");

			resultBuffer.append("Number of events:" + event_count + "\n");
			resultBuffer.append("Threat Level:" + threat + "\n");
			resultBuffer.append("Impact Level:" + impact + "\n");

			//

			fout.write(resultBuffer.toString().getBytes());

			fout.flush();
			fout.close();

		} catch (Exception e) {

			e.printStackTrace();

		}

		// reverting to snap shot

		endPointMachine.revertToCurrentSnapshot_Task(null);

		Thread.sleep(480000);

	}

	public static void pauseEndpoint(String eid) throws Exception {

		String pause = "{ " + "  \"type\":\"pause_many\", "
				+ "  \"endpoints\":[ " + "\"" + eid + "\" " + " ] " + "} ";

		HttpClientUtils.composePostRequest(
				Configuration.getProperty("restapirooturi")
						+ "/endpoints/command", pause);

	}

	public static void resumeEndpoint(String eid) throws Exception {

		String resume = "{ " + "  \"type\":\"resume_many\", "
				+ "  \"endpoints\":[ " + "    \"" + eid + "\" " + "] " + "} ";

		HttpClientUtils.composePostRequest(
				Configuration.getProperty("restapirooturi")
						+ "/endpoints/command", resume);

	}

	public static String getDriverVersion(String eid, ImpalaConnector iConnector) {

		String driver_Version = null;
		ResultSet endpointStatRes = iConnector
				.executeQueryAndReturnResults("select driver_version from endpoints where id='"
						+ eid + "'");
		try {
			while (endpointStatRes.next()) {
				try {
					driver_Version = endpointStatRes.getString(1);
				} catch (Exception e) {

				}

				System.out.println("Driver version>>>>>>" + driver_Version);

			}
		} catch (Exception e) {

		}

		return driver_Version;

	}

}
